package com.redis.sample.rate.limit;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.redisson.api.BatchResult;
import org.redisson.api.RBatch;
import org.redisson.api.RedissonClient;
import org.springframework.stereotype.Component;

@Slf4j
@Aspect
@Component
@RequiredArgsConstructor
public class RateLimitAspect {

  private final RedissonClient redissonClient;

  @Around("@annotation(rateLimit)")
  public Object limit(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
    // 获取注解中的限流设置
    int limit = rateLimit.value();
    long duration = rateLimit.duration();
    String limitKey = rateLimit.limitKey();

    // 生成 key
    String key = generateKey(limitKey);

    // 根据限流策略类型，执行相应的限流策略
    Integer cnt = count(duration, key);

    log.info("cnt = {}", cnt);
    if(cnt > limit){
      Map<String, Object> map = new HashMap<>();
      map.put("limit", "Rate limit exceeded");
      return map;
    }

    // 如果没有超过限流阈值，执行方法
    return joinPoint.proceed();
  }

  private String generateKey(String limitKey) {
    // 根据限流维度生成不同的 key
    return "rate_limit:" + limitKey;
  }

  /**
   * 统计[windowTime-expireTime, windowTime]区间的请求次数，windowTime即当前时间
   *
   * @param expireTime 窗口过期时间，单位：秒
   * @param key        redis key
   * @return 请求次数
   */
  public Integer count(long expireTime, String key) {
    // 窗口结束时间
    long windowTime = System.currentTimeMillis();
    // 窗口开始时间
    long windowStartTime = windowTime - expireTime * 1000L;

    try {
      // 创建 RBatch 实例，批量执行命令
      RBatch batch = redissonClient.createBatch();

      // 添加元素 score=当前时间戳 value=请求序列号，唯一不可重复
      batch.getScoredSortedSet(key).addAsync(windowTime, UUID.randomUUID().toString());
      // 统计[windowTime-expireTime, windowTime]区间的数据
      batch.getScoredSortedSet(key).countAsync(windowStartTime, true, windowTime, true);
      // 清除窗口过期成员
      batch.getScoredSortedSet(key).removeRangeByScoreAsync(0, true, windowStartTime, false);
      // 设置 key 过期时间
      batch.getScoredSortedSet(key).expireAsync(expireTime, TimeUnit.SECONDS);

      // 执行管道命令
      BatchResult<?> batchResult = batch.execute();

      // 返回统计数量
      List<?> responses = batchResult.getResponses();
      return (Integer) responses.get(1);
    } catch (Exception e) {
      log.error("统计请求次数异常！", e);
      return null;
    }
  }

}

